//
//  word2vec.cpp
//  word2vec
//
//  Created by LazyLie on 14-8-11.
//  Copyright (c) 2014年 NLP Group. All rights reserved.
//

#include "word2vec.h"
#include <iostream>
#include <fstream>
#include <cmath>
#include <ctime>
#include <thread>

word2vec::word2vec() {
    vec_size = 50;
    alpha_start = 0.025;
    vocabulary_size = 0;
    window_size = 5;
    train_actual = 0;
}

void word2vec::init(std::string filename) {
    std::cout << "word2vec::init learn vocabulary" << std::endl;
    std::ifstream reader(filename);
    if (!reader.is_open()) {
        std::cout << "word2vec::init Fail: can't open file " << std::endl;
        exit(0);
    }
    
    std::string word;
    std::map<std::string, int> word_counter;
    
    
    while (reader >> word) {
        if (word_counter.find(word) != word_counter.end()) {
            word_counter[word] ++;
        } else {
            word_counter[word] = 1;
        }
    }
    
    std::vector< std::pair<int, int> > counter;
    for (auto &entry : word_counter) {
        if (entry.second >= MIN_WORD_APPEARENCE) {
            vocabulary[entry.first] = vocabulary_size++;
            counter.push_back(std::make_pair(entry.second, vocabulary[entry.first]));
        }
    }
    if (vocabulary_size < 3) {
        std::cout << "word2vec::init inproper material" << std::endl;
        exit(0);
    }
    std::cout << "word2vec::init vocabulary size: " << vocabulary_size << std::endl;
    
    
    std::cout << "word2vec::init initialize huffman tree" << std::endl;
    tree.init(counter);

    
    std::cout << "word2vec::init load material" << std::endl;
    reader.close();
    reader.open(filename);
    if (!reader.is_open()) {
        std::cout << "word2vec::init Fail: can't open file" << std::endl;
        exit(0);
    }
    while (reader >> word) {
        if (vocabulary.find(word) != vocabulary.end()) {
            material.push_back(vocabulary[word]);
        }
    }
    std::cout << "word2vec::init material size (words): " << material.size() << std::endl;
}

void word2vec::train_thread(int st, int ed) {
    for (int i = st + window_size; i < ed - window_size; i++) {
        train_actual++;
        if (train_actual % (material.size() / 1000 + 1) == 0) {
            io_mutex.lock();
            std::cout << "word2vec::train train vectors: " << (double)train_actual / material.size() * 100 << "%" << std::endl;
            io_mutex.unlock();
        }
        double alpha = (1.0 - (double)train_actual / material.size()) * alpha_start;
        if (alpha < 0.0001 * alpha_start) {
            alpha = 0.0001 * alpha_start;
        }
        huffman::route_t route = tree.get_route(material[i]);
        vec_t vec_project, vec_modify(window_size * 2 * vec_size);
        for (int j = i - window_size; j <= i + window_size; j++) {
            if (j == i) continue;
            
            vec_word_mutex[material[j]].lock();
            vec_project.insert(vec_project.end(), vec_word[material[j]].begin(), vec_word[material[j]].end());
            vec_word_mutex[material[j]].unlock();
        }
        
        for (int j = 0; j < route.size(); j++) {
            vec_tree_mutex[route[j].first].lock();
            double rate = alpha * (1 - route[j].second -
                                   sigmoid(dot(vec_project, vec_tree[route[j].first])));
            
            vec_modify = add(vec_modify,
                             mul(rate, vec_tree[route[j].first]));

            vec_tree[route[j].first] = add(vec_tree[route[j].first],
                                           mul(rate, vec_project));

            vec_tree_mutex[route[j].first].unlock();
        }
        
        for (int j = i - window_size, k = 0; j <= i + window_size; j++) {
            if (i == j) continue;
            vec_word_mutex[material[j]].lock();
            vec_word[material[j]] = add(vec_word[material[j]],
                                        vec_t(vec_modify.begin() + k, vec_modify.begin() + k + vec_size));
            vec_word_mutex[material[j]].unlock();
            k += vec_size;
        }
    }
}

void word2vec::train() {
    std::cout << "word2vec::train initialize vectors" << std::endl;
    vec_word.resize(vocabulary_size, vec_t(vec_size));
    vec_tree.resize(vocabulary_size - 1, vec_t(window_size * 2 * vec_size));
    srand((unsigned int)time(0));
    for (int i = 0; i < vec_size; i++) {
        for (int j = 0; j < vocabulary_size; j++) {
            vec_word[j][i] = ((double)rand() / RAND_MAX - 0.5) / vec_size;
        }
    }
    
    
    std::cout << "word2vec::train train vectors" << std::endl;
    
    std::thread p[THREAD_NUM];
    int size = material.size() / THREAD_NUM;

    vec_word_mutex = std::vector<std::mutex>(vec_word.size());
    vec_tree_mutex = std::vector<std::mutex>(vec_tree.size());
  
    for (int i = 0; i < THREAD_NUM; i++) {
        p[i] = std::thread(std::bind(&word2vec::train_thread, this, i * size, i * size + size));
    }

    for (int i = 0; i < THREAD_NUM; i++) {
        p[i].join();
    }

    std::cout << "Done!" << std::endl;
}

word2vec::result_t word2vec::get_result() {
    result_t result;
    for (auto &entry : vocabulary) {
        result.push_back(std::make_pair(entry.first, vec_word[entry.second]));
    }
    return result;
}

double word2vec::sigmoid(double x) {
    return 1.0 / (1.0 + exp(-x));
}

double word2vec::dot(const vec_t &a, const vec_t &b) {
    if (a.size() != b.size()) {
        std::cout << "Fail: dot receives two vectors with different size. " << std::endl;
        exit(0);
    }
    double c = 0;
    for (int i = 0; i < a.size(); i++)
        c += a[i] * b[i];
    return c;
}

word2vec::vec_t word2vec::mul(double a, const vec_t &b) {
    vec_t c(b.size());
    for (int i = 0; i < b.size(); i++)
        c[i] = a * b[i];
    return c;
}

word2vec::vec_t word2vec::add(const vec_t &a, const vec_t &b) {
    if (a.size() != b.size()) {
        std::cout << "Fail: add receives two vectors with different size. " << std::endl;
        exit(0);
    }
    vec_t c(a.size());
    for (int i = 0; i < a.size(); i++)
        c[i] = a[i] + b[i];
    return c;
}